#pragma once

#include <memory>
#include <vulkan/vulkan.h>
#include "vkinfer/runtime/vk/vkblob.h"

namespace vkinfer
{
    class VkCompute
    {
    private:
        VkDevice device;
        VkCommandPool commandPool;
        VkCommandBuffer commandBuffer;
        VkQueue queue;
    
    public:
        VkCompute(VkDevice device, VkQueue queue, uint32_t queue_index);
        ~VkCompute();

        void begin_command_buffer();
        void end_command_buffer();

        void submit(VkFence fence);
        
        void dispatch(uint32_t gx, uint32_t gy, uint32_t gz)
        {
            vkCmdDispatch(commandBuffer, gx, gy, gz);
        }
        
        void bind_pipeline(VkPipeline& pipeline)
        {
            vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, pipeline);
        }

        void bind_descriptor_set(VkPipelineLayout layout, VkDescriptorSet set)
        {
            vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, layout,
                0, 1, &set, 0, NULL);
        }

        template<typename PCType>
        void push_constants(const PCType& pushconstants, VkPipelineLayout layout)
        {
            vkCmdPushConstants(commandBuffer, layout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(PCType), &pushconstants);
        }

        void buffer_barrier(const VkDependencyInfo& dep);
        void barrier_writeread(const std::vector<std::shared_ptr<VkBlob>>& buffers, bool top_layer);
    };
}
